from cryptography import x509
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.hashes import SHA256
from cryptography.x509.oid import NameOID
import datetime

# 生成 RSA 密钥对
private_key = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048,
    backend=default_backend()
)

# 生成证书请求
subject = issuer = x509.Name([
    x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
    x509.NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),
    x509.NameAttribute(NameOID.ORGANIZATION_NAME, u"My Organization"),
    x509.NameAttribute(NameOID.COMMON_NAME, u"127.0.0.1"),
])

cert_builder = x509.CertificateBuilder()
cert_builder = cert_builder.subject_name(subject)
cert_builder = cert_builder.issuer_name(issuer)
cert_builder = cert_builder.not_valid_before(datetime.datetime.utcnow())
cert_builder = cert_builder.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
cert_builder = cert_builder.serial_number(x509.random_serial_number())
cert_builder = cert_builder.public_key(private_key.public_key())

# 添加扩展（可选）
cert_builder = cert_builder.add_extension(
    x509.SubjectAlternativeName([x509.DNSName(u"127.0.0.1")]),
    critical=False,
)

# 使用 SHA256 对证书请求进行签名
certificate = cert_builder.sign(
    private_key=private_key,
    algorithm=SHA256(),
    backend=default_backend()
)

# 将私钥和证书导出为 PEM 格式
private_key_pem = private_key.private_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PrivateFormat.TraditionalOpenSSL,
    encryption_algorithm=serialization.NoEncryption()
)

certificate_pem = certificate.public_bytes(
    encoding=serialization.Encoding.PEM
)

# 将私钥和证书保存到文件中
with open("../filesOutput/private_key.pem", "wb") as key_file:
    key_file.write(private_key_pem)

with open("../filesOutput/certificate.pem", "wb") as cert_file:
    cert_file.write(certificate_pem)

print("自签名证书和私钥已生成并保存到 private_key.pem 和 certificate.pem")